#ifndef __CUDA_HOOK_HOOK_H__
#define __CUDA_HOOK_HOOK_H__


/** 
    dlfcn.h 头文件是C/C++标准库中的一个头文件，提供了用于动态链接库（shared library）的函数和数据类型，通常用于动态加
    载和卸载共享库。它主要用于实现动态链接和运行时加载共享库的功能。

    以下是 dlfcn.h 头文件的主要功能和一些相关函数和数据类型：

    - dlopen 函数：

        void *dlopen(const char *filename, int flag)：用于打开共享库文件，返回一个句柄（handle），可以用于后续的操作。
        filename 参数指定要打开的共享库文件的路径。
        flag 参数指定打开库的方式，例如 RTLD_LAZY 表示延迟加载共享库中的符号，RTLD_NOW 表示立即加载。

    - dlsym 函数：

        void *dlsym(void *handle, const char *symbol)：用于在打开的共享库中查找指定名称的符号（函数或变量）的地址。
        handle 参数是通过 dlopen 返回的句柄。
        symbol 参数是要查找的符号的名称。

    - dlclose 函数：

        int dlclose(void *handle)：用于关闭已打开的共享库，释放相关资源。
        handle 参数是通过 dlopen 返回的句柄。

    - dlerror 函数：

        const char *dlerror(void)：用于获取最近一次发生的动态链接错误的描述。
        如果在调用其他 dl* 函数时出现错误，可以使用 dlerror 来获取错误信息。

    数据类型：

        void *：在 dlopen 和 dlsym 中用于表示共享库句柄和符号地址的通用指针类型。
        RTLD_LAZY 和 RTLD_NOW：用于指定共享库加载方式的标志。

        dlfcn.h 头文件允许程序在运行时动态加载共享库，并通过函数指针或符号名称调用共享库中的函数或访问共享库中的变量。这对于实现
    插件系统、动态加载插件或实现模块化程序非常有用。

        需要注意的是，dlfcn.h 是标准C库的一部分，但在不同的操作系统上，动态链接库的加载和使用方式可能有所不同。因此，在使用动态
    链接库时，需要确保程序在目标平台上具有适当的可移植性。
*/
#include <dlfcn.h>
#include <string>
#include "macro.h"

// void *dlopen(const char *filename, int flag)
#define HOOK_GET_SYMBOL(type, symbol_name)                                          \
    do {                                                                            \
        static void *type##_handle = dlopen(s_##type##_dso, RTLD_NOW | RTLD_LOCAL); \
        HOOK_CHECK(type##_handle);                                                  \
        return dlsym(type##_handle, symbol_name.c_str());                           \
    } while (0)


class Hook {
public:
    Hook()  = default;
    ~Hook() = default;
public:    
    static void *GetCUDASymbol(const std::string &symbol_name) {
        HOOK_GET_SYMBOL(cuda, symbol_name);
    }

    // static void *GetNVMLSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(nvml, symbol_name);
    // }

    static void *GetCUDARTSymbol(const std::string &symbol_name) {
        HOOK_GET_SYMBOL(cudart, symbol_name);
    }

    // static void *GetCUDNNSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(cudnn, symbol_name);
    // }

    static void *GetCUBLASSymbol(const std::string &symbol_name) {
        HOOK_GET_SYMBOL(cublas, symbol_name);
    }

    // static void *GetCUBLASLTSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(cublasLt, symbol_name);
    // }

    // static void *GetCUFFTSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(cufft, symbol_name);
    // }

    // static void *GetNVTXSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(nvtx, symbol_name);
    // }

    // static void *GetNVRTCSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(nvrtc, symbol_name);
    // }

    // static void *GetCURANDSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(curand, symbol_name);
    // }

    // static void *GetCUSPARSESymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(cusparse, symbol_name);
    // }

    // static void *GetCUSOLVERSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(cusolver, symbol_name);
    // }

    // static void *GetNVJPEGSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(nvjpeg, symbol_name);
    // }

    // static void *GetNVBLASSymbol(const std::string &symbol_name) {
    //     HOOK_GET_SYMBOL(nvblas, symbol_name);
    // }
private:
    // nvidia native cuda dynamic library can be modified to any other unambiguous name, or moved to any path
    static constexpr const char *s_cuda_dso     = "/usr/lib/x86_64-linux-gnu/libcuda.so";
    // static constexpr const char *s_nvml_dso     = "/usr/lib/x86_64-linux-gnu/libnvidia-ml.so";
    static constexpr const char *s_cudart_dso   = "/usr/local/cuda/targets/x86_64-linux/lib/libcudart.so";
    // static constexpr const char *s_cudnn_dso    = "/usr/local/cudnn/lib64/libcudnn.so";
    static constexpr const char *s_cublas_dso   = "/usr/local/cuda/targets/x86_64-linux/lib/libcublas.so";
    // static constexpr const char *s_cublasLt_dso = "/usr/local/cuda/targets/x86_64-linux/lib/libcublasLt.so";
    // static constexpr const char *s_cufft_dso    = "/usr/local/cuda/targets/x86_64-linux/lib/libcufft.so";
    // static constexpr const char *s_nvtx_dso     = "/usr/local/cuda/targets/x86_64-linux/lib/libnvToolsExt.so";
    // static constexpr const char *s_nvrtc_dso    = "/usr/local/cuda/targets/x86_64-linux/lib/libnvrtc.so";
    // static constexpr const char *s_curand_dso   = "/usr/local/cuda/targets/x86_64-linux/lib/libcurand.so";
    // static constexpr const char *s_cusparse_dso = "/usr/local/cuda/targets/x86_64-linux/lib/libcusparse.so";
    // static constexpr const char *s_cusolver_dso = "/usr/local/cuda/targets/x86_64-linux/lib/libcusolver.so";
    // static constexpr const char *s_nvjpeg_dso   = "/usr/local/cuda/targets/x86_64-linux/lib/libnvjpeg.so";
    // static constexpr const char *s_nvblas_dso   = "/usr/local/cuda/targets/x86_64-linux/lib/libnvblas.so";
    HOOK_DISABLE_COPY_AND_ASSIGN(Hook)
};

#define HOOK_CUDA_SYMBOL(symbol_name)     Hook::GetCUDASymbol(symbol_name)
// #define HOOK_NVML_SYMBOL(symbol_name)     Hook::GetNVMLSymbol(symbol_name)
#define HOOK_CUDART_SYMBOL(symbol_name)   Hook::GetCUDARTSymbol(symbol_name)
// #define HOOK_CUDNN_SYMBOL(symbol_name)    Hook::GetCUDNNSymbol(symbol_name)
#define HOOK_CUBLAS_SYMBOL(symbol_name)   Hook::GetCUBLASSymbol(symbol_name)
// #define HOOK_CUBLASLT_SYMBOL(symbol_name) Hook::GetCUBLASLTSymbol(symbol_name)
// #define HOOK_CUFFT_SYMBOL(symbol_name)    Hook::GetCUFFTSymbol(symbol_name)
// #define HOOK_NVTX_SYMBOL(symbol_name)     Hook::GetNVTXSymbol(symbol_name)
// #define HOOK_NVRTC_SYMBOL(symbol_name)    Hook::GetNVRTCSymbol(symbol_name)
// #define HOOK_CURAND_SYMBOL(symbol_name)   Hook::GetCURANDSymbol(symbol_name)
// #define HOOK_CUSPARSE_SYMBOL(symbol_name) Hook::GetCUSPARSESymbol(symbol_name)
// #define HOOK_CUSOLVER_SYMBOL(symbol_name) Hook::GetCUSOLVERSymbol(symbol_name)
// #define HOOK_NVJPEG_SYMBOL(symbol_name)   Hook::GetNVJPEGSymbol(symbol_name)
// #define HOOK_NVBLAS_SYMBOL(symbol_name)   Hook::GetNVBLASSymbol(symbol_name)

#endif // __CUDA_HOOK_HOOK_H__